#!/usr/bin/env python3
"""
Network Discovery Tool for EchoKitBot System
自动发现网络中的设备IP地址

功能特性：
- 自动检测本地网络子网
- 优先扫描常用IP地址
- 多端口检测支持
- 快速超时机制
- 并发扫描提高速度
"""

import socket
import subprocess
import ipaddress
import concurrent.futures
from typing import Optional, List, Tuple
import time


class NetworkDiscovery:
    """网络设备自动发现工具"""
    
    # 常用IP优先级列表（最后一位）
    PRIORITY_IPS = [28, 35, 42, 40, 50, 100, 1, 10, 20, 30, 254]
    
    # 默认超时时间（秒）
    DEFAULT_TIMEOUT = 0.5
    
    # 最大并发扫描数
    MAX_WORKERS = 20
    
    def __init__(self, timeout: float = DEFAULT_TIMEOUT, verbose: bool = False):
        """
        初始化网络发现工具
        
        Args:
            timeout: 连接超时时间（秒）
            verbose: 是否输出详细日志
        """
        self.timeout = timeout
        self.verbose = verbose
    
    def get_local_ip(self) -> Optional[str]:
        """
        获取本机IP地址
        
        Returns:
            本机IP地址字符串，失败返回None
        """
        try:
            # 创建UDP socket连接到外部地址（不实际发送数据）
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            local_ip = s.getsockname()[0]
            s.close()
            return local_ip
        except Exception as e:
            if self.verbose:
                print(f"获取本机IP失败: {e}")
            return None
    
    def get_subnet(self, ip: str, prefix: int = 24) -> Optional[str]:
        """
        从IP地址获取子网
        
        Args:
            ip: IP地址
            prefix: 子网前缀长度（默认24，即255.255.255.0）
        
        Returns:
            子网字符串（如 "<LAN_SUBNET>"），失败返回None
        """
        try:
            network = ipaddress.IPv4Network(f"{ip}/{prefix}", strict=False)
            return str(network)
        except Exception as e:
            if self.verbose:
                print(f"获取子网失败: {e}")
            return None
    
    def check_port(self, ip: str, port: int) -> bool:
        """
        检查指定IP和端口是否可达
        
        Args:
            ip: 目标IP地址
            port: 目标端口
        
        Returns:
            端口开放返回True，否则返回False
        """
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.timeout)
            result = sock.connect_ex((ip, port))
            sock.close()
            return result == 0
        except Exception:
            return False
    
    def check_host_alive(self, ip: str) -> bool:
        """
        使用ping检查主机是否存活
        
        Args:
            ip: 目标IP地址
        
        Returns:
            主机存活返回True，否则返回False
        """
        try:
            # Linux/macOS使用 -c 1 -W 1
            result = subprocess.run(
                ['ping', '-c', '1', '-W', '1', ip],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=2
            )
            return result.returncode == 0
        except Exception:
            return False
    
    def scan_ip_port(self, ip: str, ports: List[int]) -> Tuple[str, List[int]]:
        """
        扫描单个IP的多个端口
        
        Args:
            ip: 目标IP地址
            ports: 端口列表
        
        Returns:
            (IP地址, 开放的端口列表)
        """
        open_ports = []
        for port in ports:
            if self.check_port(ip, port):
                open_ports.append(port)
                if self.verbose:
                    print(f"✓ 发现设备: {ip}:{port}")
        return ip, open_ports
    
    def discover_device(
        self,
        ports: List[int],
        subnet: Optional[str] = None,
        priority_only: bool = True
    ) -> Optional[str]:
        """
        自动发现网络中的设备
        
        Args:
            ports: 要检测的端口列表
            subnet: 子网（如 "192.168.2.0/24"），None则自动检测
            priority_only: 是否仅扫描优先IP列表
        
        Returns:
            找到的第一个设备IP地址，未找到返回None
        """
        # 获取子网
        if subnet is None:
            local_ip = self.get_local_ip()
            if local_ip is None:
                if self.verbose:
                    print("无法获取本机IP地址")
                return None
            subnet = self.get_subnet(local_ip)
            if subnet is None:
                if self.verbose:
                    print("无法获取子网信息")
                return None
        
        if self.verbose:
            print(f"扫描子网: {subnet}")
            print(f"检测端口: {ports}")
        
        # 解析子网
        try:
            network = ipaddress.IPv4Network(subnet, strict=False)
            base_ip = str(network.network_address)
            prefix = base_ip.rsplit('.', 1)[0]
        except Exception as e:
            if self.verbose:
                print(f"解析子网失败: {e}")
            return None
        
        # 生成要扫描的IP列表
        if priority_only:
            # 仅扫描优先IP
            ips_to_scan = [f"{prefix}.{last}" for last in self.PRIORITY_IPS]
        else:
            # 扫描整个子网（优先IP在前）
            priority_ips = [f"{prefix}.{last}" for last in self.PRIORITY_IPS]
            all_ips = [str(ip) for ip in network.hosts()]
            # 移除已在优先列表中的IP
            other_ips = [ip for ip in all_ips if ip not in priority_ips]
            ips_to_scan = priority_ips + other_ips
        
        if self.verbose:
            print(f"扫描 {len(ips_to_scan)} 个IP地址...")
        
        # 并发扫描
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
            future_to_ip = {
                executor.submit(self.scan_ip_port, ip, ports): ip 
                for ip in ips_to_scan
            }
            
            for future in concurrent.futures.as_completed(future_to_ip):
                ip, open_ports = future.result()
                if open_ports:
                    # 找到第一个有开放端口的设备就返回
                    if self.verbose:
                        print(f"✓ 找到设备: {ip} (端口: {open_ports})")
                    return ip
        
        if self.verbose:
            print("未找到任何设备")
        return None
    
    def discover_echokit_server(
        self,
        default_ip: str = "<ECHO_SERVER>",
        ports: Optional[List[int]] = None
    ) -> str:
        """
        发现EchoKit服务器IP地址
        
        Args:
            default_ip: 默认IP地址（发现失败时使用）
            ports: 要检测的端口列表（默认[8060]）
        
        Returns:
            EchoKit服务器IP地址
        """
        if ports is None:
            ports = [8060]
        
        if self.verbose:
            print("正在发现EchoKit服务器...")
        
        discovered_ip = self.discover_device(ports, priority_only=True)
        
        if discovered_ip:
            if self.verbose:
                print(f"✓ 自动发现EchoKit服务器: {discovered_ip}")
            return discovered_ip
        else:
            if self.verbose:
                print(f"⚠ 未发现EchoKit服务器，使用默认IP: {default_ip}")
            return default_ip
    
    def discover_robot(
        self,
        default_ip: str = "<ECHOKIT_IP>",
        ports: Optional[List[int]] = None
    ) -> str:
        """
        发现机器人IP地址
        
        Args:
            default_ip: 默认IP地址（发现失败时使用）
            ports: 要检测的端口列表（默认使用ping检测）
        
        Returns:
            机器人IP地址
        """
        if self.verbose:
            print("正在发现机器人...")
        
        # 如果指定了端口，使用端口扫描
        if ports:
            discovered_ip = self.discover_device(ports, priority_only=True)
            if discovered_ip:
                if self.verbose:
                    print(f"✓ 自动发现机器人: {discovered_ip}")
                return discovered_ip
        
        # 否则使用ping扫描优先IP列表
        local_ip = self.get_local_ip()
        if local_ip:
            prefix = local_ip.rsplit('.', 1)[0]
            for last in self.PRIORITY_IPS:
                ip = f"{prefix}.{last}"
                if self.check_host_alive(ip):
                    if self.verbose:
                        print(f"✓ 自动发现机器人: {ip}")
                    return ip
        
        if self.verbose:
            print(f"⚠ 未发现机器人，使用默认IP: {default_ip}")
        return default_ip
    
    def discover_camera(
        self,
        default_ip: str = "<CAMERA_IP>",
        ports: Optional[List[int]] = None
    ) -> str:
        """
        发现摄像头IP地址
        
        Args:
            default_ip: 默认IP地址（发现失败时使用）
            ports: 要检测的端口列表（默认[8554]）
        
        Returns:
            摄像头IP地址
        """
        if ports is None:
            ports = [8554]  # RTSP默认端口
        
        if self.verbose:
            print("正在发现摄像头...")
        
        discovered_ip = self.discover_device(ports, priority_only=True)
        
        if discovered_ip:
            if self.verbose:
                print(f"✓ 自动发现摄像头: {discovered_ip}")
            return discovered_ip
        else:
            if self.verbose:
                print(f"⚠ 未发现摄像头，使用默认IP: {default_ip}")
            return default_ip


def quick_discover_echokit(
    default_ip: str = "<ECHO_SERVER>",
    verbose: bool = False
) -> str:
    """
    快速发现EchoKit服务器（便捷函数）
    
    Args:
        default_ip: 默认IP地址
        verbose: 是否输出详细日志
    
    Returns:
        EchoKit服务器IP地址
    """
    discovery = NetworkDiscovery(timeout=0.3, verbose=verbose)
    return discovery.discover_echokit_server(default_ip)


def quick_discover_robot(
    default_ip: str = "<ECHOKIT_IP>",
    verbose: bool = False
) -> str:
    """
    快速发现机器人（便捷函数）
    
    Args:
        default_ip: 默认IP地址
        verbose: 是否输出详细日志
    
    Returns:
        机器人IP地址
    """
    discovery = NetworkDiscovery(timeout=0.3, verbose=verbose)
    return discovery.discover_robot(default_ip)


def quick_discover_camera(
    default_ip: str = "<CAMERA_IP>",
    verbose: bool = False
) -> str:
    """
    快速发现摄像头（便捷函数）
    
    Args:
        default_ip: 默认IP地址
        verbose: 是否输出详细日志
    
    Returns:
        摄像头IP地址
    """
    discovery = NetworkDiscovery(timeout=0.3, verbose=verbose)
    return discovery.discover_camera(default_ip)


# 命令行测试
if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='网络设备自动发现工具')
    parser.add_argument('--device', choices=['echokit', 'robot', 'camera', 'all'],
                       default='all', help='要发现的设备类型')
    parser.add_argument('--verbose', '-v', action='store_true',
                       help='显示详细输出')
    parser.add_argument('--full-scan', action='store_true',
                       help='扫描整个子网（而非仅优先IP）')
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("EchoKitBot 网络设备自动发现工具")
    print("=" * 60)
    print()
    
    discovery = NetworkDiscovery(timeout=0.5, verbose=args.verbose)
    
    # 显示本机网络信息
    local_ip = discovery.get_local_ip()
    if local_ip:
        print(f"本机IP: {local_ip}")
        subnet = discovery.get_subnet(local_ip)
        if subnet:
            print(f"子网: {subnet}")
    print()
    
    # 发现设备
    if args.device in ['echokit', 'all']:
        print("🔍 发现EchoKit服务器...")
        start_time = time.time()
        echokit_ip = discovery.discover_echokit_server()
        elapsed = time.time() - start_time
        print(f"✓ EchoKit服务器: {echokit_ip} (耗时: {elapsed:.2f}秒)")
        print()
    
    if args.device in ['robot', 'all']:
        print("🔍 发现机器人...")
        start_time = time.time()
        robot_ip = discovery.discover_robot()
        elapsed = time.time() - start_time
        print(f"✓ 机器人: {robot_ip} (耗时: {elapsed:.2f}秒)")
        print()
    
    if args.device in ['camera', 'all']:
        print("🔍 发现摄像头...")
        start_time = time.time()
        camera_ip = discovery.discover_camera()
        elapsed = time.time() - start_time
        print(f"✓ 摄像头: {camera_ip} (耗时: {elapsed:.2f}秒)")
        print()
    
    print("=" * 60)
    print("发现完成")
    print("=" * 60)
